"""
DeepLabv3+ 模型实现
基于ResNet-50的DeepLabv3+，使用空洞卷积和ASPP模块
注意：需要3通道输入以使用ImageNet预训练权重
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50, ResNet50_Weights


class ASPP(nn.Module):
    """
    Atrous Spatial Pyramid Pooling (ASPP) 模块
    使用不同膨胀率的空洞卷积捕获多尺度信息
    """

    def __init__(self, in_channels, out_channels=256):
        super(ASPP, self).__init__()

        # 1x1卷积
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

        # 3x3空洞卷积，rate=6
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=6, dilation=6, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

        # 3x3空洞卷积，rate=12
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=12, dilation=12, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

        # 3x3空洞卷积，rate=18
        self.conv4 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=18, dilation=18, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

        # 全局平均池化
        self.global_avg_pool = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

        # 融合后的1x1卷积
        self.conv_out = nn.Sequential(
            nn.Conv2d(out_channels * 5, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5)
        )

    def forward(self, x):
        size = x.shape[2:]

        # 各个分支
        x1 = self.conv1(x)
        x2 = self.conv2(x)
        x3 = self.conv3(x)
        x4 = self.conv4(x)
        x5 = self.global_avg_pool(x)
        x5 = F.interpolate(x5, size=size, mode='bilinear', align_corners=True)

        # 拼接
        x = torch.cat([x1, x2, x3, x4, x5], dim=1)

        # 融合
        x = self.conv_out(x)

        return x


class DeepLabV3Plus(nn.Module):
    """
    DeepLabv3+ 模型
    输入：(B, 3, H, W) - 必须是3通道
    输出：(B, out_channels, H, W)
    """

    def __init__(self, in_channels=3, out_channels=1, pretrained=True):
        """
        Args:
            in_channels: 输入通道数（必须是3，用于ImageNet预训练）
            out_channels: 输出通道数
            pretrained: 是否使用预训练权重
        """
        super(DeepLabV3Plus, self).__init__()

        if in_channels != 3:
            print(f"[WARNING] DeepLabv3+ 需要3通道输入以使用预训练权重，当前输入通道数: {in_channels}")

        # 使用ResNet-50作为backbone
        if pretrained:
            resnet = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
        else:
            resnet = resnet50(weights=None)

        # 编码器（ResNet的前几层）
        self.conv1 = resnet.conv1
        self.bn1 = resnet.bn1
        self.relu = resnet.relu
        self.maxpool = resnet.maxpool

        self.layer1 = resnet.layer1  # 256 channels
        self.layer2 = resnet.layer2  # 512 channels
        self.layer3 = resnet.layer3  # 1024 channels
        self.layer4 = resnet.layer4  # 2048 channels

        # ASPP模块
        self.aspp = ASPP(in_channels=2048, out_channels=256)

        # 低层特征处理（来自layer1）
        self.low_level_conv = nn.Sequential(
            nn.Conv2d(256, 48, 1, bias=False),
            nn.BatchNorm2d(48),
            nn.ReLU(inplace=True)
        )

        # 解码器
        self.decoder = nn.Sequential(
            nn.Conv2d(256 + 48, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Conv2d(256, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1)
        )

        # 输出层
        self.classifier = nn.Conv2d(256, out_channels, 1)

    def forward(self, x):
        size = x.shape[2:]

        # 编码器
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        low_level_feat = self.layer1(x)  # 1/4
        x = self.layer2(low_level_feat)  # 1/8
        x = self.layer3(x)  # 1/16
        x = self.layer4(x)  # 1/32

        # ASPP
        x = self.aspp(x)

        # 上采样到1/4
        x = F.interpolate(x, size=low_level_feat.shape[2:], mode='bilinear', align_corners=True)

        # 处理低层特征
        low_level_feat = self.low_level_conv(low_level_feat)

        # 拼接
        x = torch.cat([x, low_level_feat], dim=1)

        # 解码
        x = self.decoder(x)

        # 分类
        x = self.classifier(x)

        # 上采样到原始尺寸
        x = F.interpolate(x, size=size, mode='bilinear', align_corners=True)

        return x


if __name__ == '__main__':
    # 测试模型
    model = DeepLabV3Plus(in_channels=3, out_channels=1, pretrained=False)
    x = torch.randn(2, 3, 256, 256)
    y = model(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {y.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
